from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from PIL import Image
import time
import io

chrome_options = Options()
chrome_options.add_argument("--headless")  # Run in headless mode
chrome_options.add_argument("--disable-gpu")
chrome_options.add_argument("--window-size=1920,1080")


# 启动浏览器（这里以Chrome为例）
service = Service('C:\javasoft\chromedriver-win64\chromedriver.exe')  # Replace with your chromedriver path
driver = webdriver.Chrome(service=service, options=chrome_options)
try:
    # 打开网页
    driver.get('file:///C:/projects/python/test/test.html')
    
     # 等待元素加载（根据实际情况修改定位方式和等待条件）
    element = WebDriverWait(driver, 10).until(
        EC.presence_of_element_located((By.CLASS_NAME, 'container'))
    )
    scroll_height = driver.execute_script("return arguments[0].scrollHeight;", element)
    viewport_height = element.size['height']
    num_screenshots = scroll_height // viewport_height + (1 if scroll_height % viewport_height != 0 else 0)
    print("viewport_height",viewport_height)
    print("scroll_height",scroll_height)
    print("num_screenshots",num_screenshots)
    #scroll_height = driver.execute_script("return document.getElementsByClassName('container')[0].scrollHeight")
    #js='document.getElementsByClassName("container")[0].scrollTop=10000' 
    #driver.execute_script(js)
    #driver.execute_script("arguments[0].scrollTop=100", element)

    left = element.location['x']
    top = element.location['y']
    right = element.location['x'] + element.size['width']
    bottom = element.location['y'] + element.size['height']

    stitched_image = Image.new('RGB', (element.size['width'], scroll_height))

    scroll_top = 0
    for i in range(num_screenshots):
        scroll_top = i * viewport_height
        driver.execute_script("arguments[0].scrollTop=" + str(scroll_top), element)
        screenshot_bytes = driver.get_screenshot_as_png()
        image = Image.open(io.BytesIO(screenshot_bytes))
        cropped_image = image.crop((left, top, right, bottom))
        stitched_image.paste(cropped_image, (0, scroll_top))
        
    stitched_image.save('stitched_image.png')

    
finally:
    # 关闭浏览器
    driver.quit()